#include "BasicOP.h"
#include <assert.h>
#include <float.h>
#include <math.h>

using namespace DL;

Tensor3D<float>& DL::Conv2D(const Tensor3D<float>& tensor, int stride, int padding, const Tensor4D<float>& weight, const Tensor1D<float>& bias)
{
	int outWidth = (tensor.width - weight.width + 2 * padding) / stride + 1;
	int outHeight = (tensor.height - weight.height + 2 * padding) / stride + 1;
	Tensor3D<float>* outTensor = new Tensor3D<float>(outWidth, outHeight, weight.deepth);

	for (int b = 0; b < outTensor->channel; b++)
	{
		for (int y = 0; y < outTensor->height; y++)
		{
			for (int x = 0; x < outTensor->width; x++)
			{
				float result = bias.Get(b);
				for (int c = 0; c < tensor.channel; c++)
				{
					for (int ky = 0; ky < weight.height; ky++)
					{
						for (int kx = 0; kx < weight.width; kx++)
						{
							result += tensor.Get(x * stride - padding + kx, y * stride - padding + ky, c) * weight.Get(kx, ky, c, b);
						}
					}
				}
				outTensor->Set(x, y, b, result);
			}
		}
	}
	return *outTensor;
}

Tensor3D<float>& DL::MaxPooling2D(const Tensor3D<float>& tensor, int size, int stride, int padding)
{
	int outWidth = (tensor.width - size + 2 * padding) / stride + 1;
	int outHeight = (tensor.height - size + 2 * padding) / stride + 1;
	Tensor3D<float>* outTensor = new Tensor3D<float>(outWidth, outHeight, tensor.channel);

	for (int y = 0; y < outTensor->height; y++)
	{
		for (int x = 0; x < outTensor->width; x++)
		{
			for (int c = 0; c < tensor.channel; c++)
			{
				float result = -FLT_MAX;
				for (int ky = 0; ky < size; ky++)
				{
					for (int kx = 0; kx < size; kx++)
					{
						float value = tensor.Get(x * stride - padding + kx, y * stride - padding + ky, c);
						result = result > value ? result : value;
					}
				}
				outTensor->Set(x, y, c, result);
			}
		}
	}
	return *outTensor;
}

Tensor3D<float>& DL::Relu(const Tensor3D<float>& tensor)
{
	Tensor3D<float>* outTensor = new Tensor3D<float>(tensor.width, tensor.height, tensor.channel);

	for (int y = 0; y < outTensor->height; y++)
	{
		for (int x = 0; x < outTensor->width; x++)
		{
			for (int c = 0; c < tensor.channel; c++)
			{
				float value = tensor.Get(x, y, c);
				value = value > 0 ? value : 0;
				outTensor->Set(x, y, c, value);
			}
		}
	}
	return *outTensor;
}

Tensor1D<float>& DL::Relu(const Tensor1D<float>& tensor)
{
	Tensor1D<float>* outTensor = new Tensor1D<float>(tensor.width);

	for (int x = 0; x < outTensor->width; x++)
	{
		float value = tensor.Get(x);
		value = value > 0 ? value : 0;
		outTensor->Set(x, value);
	}
	return *outTensor;
}

Tensor1D<float>& DL::Flatten(const Tensor3D<float>& tensor)
{
	Tensor1D<float>* outTensor = new Tensor1D<float>(tensor.width * tensor.height * tensor.channel);
	for (int c = 0; c < tensor.channel; c++)
	{
		for (int y = 0; y < tensor.height; y++)
		{
			for (int x = 0; x < tensor.width; x++)
			{
				float value = tensor.Get(x, y, c);
				outTensor->Set(c * tensor.framePiexlCount + y * tensor.width + x, value);
			}
		}
	}
	return *outTensor;
}

Tensor1D<float>& DL::Linear(const Tensor1D<float>& tensor, const Tensor2D<float>& weight, const Tensor1D<float>& bias)
{
	Tensor1D<float>* outTensor = new Tensor1D<float>(weight.height);
	for (int out = 0; out < outTensor->width; out++)
	{
		float result = bias.Get(out);
		for (int in = 0; in < tensor.width; in++)
		{
			result += tensor.Get(in) * weight.Get(in, out);
		}
		outTensor->Set(out, result);
	}
	return *outTensor;
}

Tensor1D<float>& DL::SoftMax(const Tensor1D<float>& tensor)
{
	Tensor1D<float>* outTensor = new Tensor1D<float>(tensor.width);
	float sum = 0;
	for (int i = 0; i < tensor.width; i++)
	{
		float ez = exp(tensor.Get(i));
		sum += ez;
		outTensor->Set(i, ez);
	}
	for (int i = 0; i < tensor.width; i++)
	{
		outTensor->Set(i, outTensor->Get(i) / sum);
	}
	return *outTensor;
}

Tensor3D<float>& DL::Normalize(const Tensor3D<float>& tensor, const float mean[], const float std[])
{
	Tensor3D<float>* outTensor = new Tensor3D<float>(tensor.width, tensor.height, tensor.channel);
	for (int c = 0; c < tensor.channel; c++)
	{
		for (int y = 0; y < tensor.height; y++)
		{
			for (int x = 0; x < tensor.width; x++)
			{
				float value = tensor.Get(x, y, c);
				value = (value - mean[c]) / std[c];
				outTensor->Set(x, y, c, value);
			}
		}
	}
	return *outTensor;
}
